package com.mezz.bones.framework.web.servlet;

import jakarta.servlet.ServletOutputStream;
import jakarta.servlet.WriteListener;
import jakarta.servlet.http.HttpServletResponse;
import jakarta.servlet.http.HttpServletResponseWrapper;

import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.UnsupportedEncodingException;

/**
 * @description: 可修改请求体的request
 * @author: lzt
 * @create: 2022-07-30 09:51
 */
public class RepeatableHttpServletResponseWrapper extends HttpServletResponseWrapper {

  private ByteArrayOutputStream byteArrayOutputStream;

  private ServletOutputStream servletOutputStream;

  public RepeatableHttpServletResponseWrapper(HttpServletResponse response) {
    super(response);
    byteArrayOutputStream = new ByteArrayOutputStream();
    servletOutputStream = new ServletOutputStream() {
      @Override
      public boolean isReady() {
        return false;
      }

      @Override
      public void setWriteListener(WriteListener writeListener) {

      }

      @Override
      public void write(int b) throws IOException {
        response.getOutputStream().write(b);
        // 同时写入字节数组
        byteArrayOutputStream.write(b);
      }
    };
  }

  @Override
  public ServletOutputStream getOutputStream() throws IOException {
    return servletOutputStream;
  }

  public byte[] toByteArray() {
    return byteArrayOutputStream.toByteArray();
  }

  public String getContent(){
    try {
      return new String(toByteArray(),"utf-8");
    } catch (UnsupportedEncodingException e) {
      throw new RuntimeException(e);
    }
  }
}
